package org.fujay.commons.langchain4j.rag.augmentor.transform;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.rag.query.Query;
import lombok.extern.slf4j.Slf4j;
import org.fujay.commons.langchain4j.rag.augmentor.EnhancedQueryTransformer;

import java.util.*;

import static dev.langchain4j.internal.Utils.getOrDefault;
import static dev.langchain4j.internal.Utils.isNotNullOrBlank;
import static dev.langchain4j.internal.ValidationUtils.ensureGreaterThanZero;
import static dev.langchain4j.internal.ValidationUtils.ensureNotNull;
import static java.util.Collections.singletonList;
import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;

/**
 * @author slm
 * @description 增强查询转换器，支持多种转换策略、中文优化和意图分析
 */
@Slf4j
public class EnhancedQueryTransformerImpl implements EnhancedQueryTransformer {

    /**
     * 默认压缩查询模板 - 针对中文优化
     */
    public static final PromptTemplate DEFAULT_COMPRESS_TEMPLATE = PromptTemplate.from(
            """
                    请阅读并理解用户与AI之间的对话历史，然后分析用户的新查询。
                    从对话历史和新查询中提取所有相关信息、关键词和上下文。
                    将该查询重新表述为清晰、简洁且独立的格式，适合信息检索使用。
                    对于中文查询，请保持语言表达的自然性和流畅性。
                    
                    对话历史:
                    {{chatMemory}}
                    
                    用户查询: {{query}}
                    
                    非常重要：你只需提供重新表述的查询，不要添加任何额外内容！
                    不要在查询前加任何前缀！""");
    /**
     * 默认扩展查询模板 - 针对中文优化
     */
    public static final PromptTemplate DEFAULT_EXPAND_TEMPLATE = PromptTemplate.from(
            """
                    请为用户的查询生成{{n}}个不同版本的查询语句。
                    每个版本应使用不同的措辞、同义词或替代句式，但都要保持原始含义。
                    这些查询版本将用于检索相关文档。
                    对于中文查询，请考虑不同的表达方式和行业术语变体。
                    
                    非常重要：每个查询版本必须单独占一行，不要添加编号、连字符或任何其他格式！
                    
                    用户查询: {{query}}""");
    /**
     * 意图分析模板 - 用于确定最佳转换策略
     */
    public static final PromptTemplate INTENT_ANALYSIS_TEMPLATE = PromptTemplate.from(
            """
                    分析用户的查询，确定最合适的查询处理策略。
                    
                    查询: {{query}}
                    
                    对话历史: {{chatMemory}}
                    
                    请从以下选项中选择一个最合适的策略，并只返回相应的选项编号（1、2或3）：
                    1. 如果查询简短明确且与对话历史无关，应保持原样
                    2. 如果查询依赖于对话历史或含有指代词，应压缩为完整查询
                    3. 如果查询表达模糊或可以有多种解释，应扩展为多个变体
                    
                    回答:""");
    /**
     * 系统提示词模板
     */
    public static final PromptTemplate SYSTEM_PROMPT_TEMPLATE = PromptTemplate.from(
            """
                    {{systemPrompt}}
                    
                    你是一个专业的查询优化助手，负责优化用户查询以提高检索效果。""");
    private final ChatModel ChatModel;
    private final TransformStrategy strategy;
    /**
     * 压缩查询提示词模板
     */
    private final PromptTemplate compressTemplate;
    /**
     * 扩展查询提示词模板
     */
    private final PromptTemplate expandTemplate;
    /**
     * 意图分析提示词模板
     */
    private final PromptTemplate intentAnalysisTemplate;
    private final int expansionCount;
    private String systemPrompt = "";

    /**
     * 创建查询转换器
     *
     * @param ChatModel 聊天语言模型
     * @param strategy          转换策略
     */
    public EnhancedQueryTransformerImpl(ChatModel ChatModel, TransformStrategy strategy) {
        this(ChatModel, strategy, DEFAULT_COMPRESS_TEMPLATE, DEFAULT_EXPAND_TEMPLATE, 3);
    }

    /**
     * 创建完整配置的查询转换器
     *
     * @param ChatModel 聊天语言模型
     * @param strategy          转换策略
     * @param compressTemplate  压缩模板
     * @param expandTemplate    扩展模板
     * @param expansionCount    扩展数量
     */
    public EnhancedQueryTransformerImpl(
            ChatModel ChatModel,
            TransformStrategy strategy,
            PromptTemplate compressTemplate,
            PromptTemplate expandTemplate,
            int expansionCount) {
        this.ChatModel = ensureNotNull(ChatModel, "ChatModel");
        this.strategy = getOrDefault(strategy, TransformStrategy.AUTO);
        this.compressTemplate = getOrDefault(compressTemplate, DEFAULT_COMPRESS_TEMPLATE);
        this.expandTemplate = getOrDefault(expandTemplate, DEFAULT_EXPAND_TEMPLATE);
        this.intentAnalysisTemplate = INTENT_ANALYSIS_TEMPLATE;
        this.expansionCount = ensureGreaterThanZero(getOrDefault(expansionCount, 3), "expansionCount");
    }

    /**
     * 创建构建器
     */
    public static EnhancedQueryTransformerBuilder builder() {
        return new EnhancedQueryTransformerBuilder();
    }

    @Override
    public Collection<Query> transform(Query query) {
        if (query == null || query.text() == null || query.text().trim().isEmpty()) {
            return Collections.emptyList();
        }

        log.debug("使用策略 {} 转换查询: {}", strategy, query.text());

        // 根据策略选择转换方法
        switch (determineStrategy(query)) {
            case ORIGINAL:
                return singletonList(query);

            case COMPRESS:
                return compressQuery(query);

            case EXPAND:
                return expandQuery(query);

            case HYBRID:
                // 先压缩再扩展
                Collection<Query> compressed = compressQuery(query);
                if (compressed.isEmpty()) {
                    return singletonList(query);
                }
                return expandQuery(compressed.iterator().next());

            default:
                log.warn("未知策略或策略选择失败，使用原始查询");
                return singletonList(query);
        }
    }

    @Override
    public String transform(String queryText) {
        Query query = Query.from(queryText);
        Collection<Query> transformed = transform(query);

        if (transformed.isEmpty()) {
            return queryText;
        }

        // 对于文本查询，仅返回第一个转换结果
        return transformed.iterator().next().text();
    }

    @Override
    public void setSystemPrompt(String systemPrompt) {
        this.systemPrompt = systemPrompt;
    }

    /**
     * 确定实际使用的转换策略
     */
    private TransformStrategy determineStrategy(Query query) {
        // 如果不是AUTO策略，直接返回设置的策略
        if (strategy != TransformStrategy.AUTO) {
            return strategy;
        }

        // 使用意图分析选择策略
        try {
            List<ChatMessage> chatMemory = query.metadata() != null ? query.metadata().chatMemory()
                    : Collections.emptyList();

            Map<String, Object> variables = new HashMap<>();
            variables.put("query", query.text());
            variables.put("chatMemory", formatChatMemory(chatMemory));

            Prompt prompt = intentAnalysisTemplate.apply(variables);
            String response = ChatModel.chat(prompt.text()).trim();

            // 解析响应，获取策略选择
            if (response.contains("1")) {
                return TransformStrategy.ORIGINAL;
            } else if (response.contains("2")) {
                return TransformStrategy.COMPRESS;
            } else if (response.contains("3")) {
                return TransformStrategy.EXPAND;
            }

            // 默认使用压缩策略
            return TransformStrategy.COMPRESS;
        } catch (Exception e) {
            log.error("意图分析失败，使用默认压缩策略", e);
            return TransformStrategy.COMPRESS;
        }
    }

    /**
     * 压缩查询
     */
    private Collection<Query> compressQuery(Query query) {
        List<ChatMessage> chatMemory = query.metadata() != null ? query.metadata().chatMemory()
                : Collections.emptyList();

        // 如果没有对话历史，无需压缩
        if (chatMemory.isEmpty()) {
            return singletonList(query);
        }

        try {
            Map<String, Object> variables = new HashMap<>();
            variables.put("query", query.text());
            variables.put("chatMemory", formatChatMemory(chatMemory));

            // 添加系统提示（如果有）
            Prompt prompt;
            if (isNotNullOrBlank(systemPrompt)) {
                Map<String, Object> systemVars = new HashMap<>();
                systemVars.put("systemPrompt", systemPrompt);
                String enhancedSystemPrompt = SYSTEM_PROMPT_TEMPLATE.apply(systemVars).text();

                variables.put("systemPrompt", enhancedSystemPrompt);
            }
            prompt = compressTemplate.apply(variables);

            String compressedQueryText = ChatModel.chat(prompt.text());
            Query compressedQuery = query.metadata() == null
                    ? Query.from(compressedQueryText)
                    : Query.from(compressedQueryText, query.metadata());

            log.debug("查询压缩: {} -> {}", query.text(), compressedQueryText);
            return singletonList(compressedQuery);
        } catch (Exception e) {
            log.error("查询压缩失败", e);
            return singletonList(query);
        }
    }

    /**
     * 扩展查询
     */
    private Collection<Query> expandQuery(Query query) {
        try {
            Map<String, Object> variables = new HashMap<>();
            variables.put("query", query.text());
            variables.put("n", expansionCount);

            // 添加系统提示（如果有）
            Prompt prompt;
            if (isNotNullOrBlank(systemPrompt)) {
                Map<String, Object> systemVars = new HashMap<>();
                systemVars.put("systemPrompt", systemPrompt);
                String enhancedSystemPrompt = SYSTEM_PROMPT_TEMPLATE.apply(systemVars).text();

                variables.put("systemPrompt", enhancedSystemPrompt);
            }
            prompt = expandTemplate.apply(variables);

            String response = ChatModel.chat(prompt.text());
            List<String> expandedQueries = parseQueryLines(response);

            // 确保至少有一个查询
            if (expandedQueries.isEmpty()) {
                return singletonList(query);
            }

            log.debug("查询扩展: {} -> {} 个查询变体", query.text(), expandedQueries.size());

            // 转换为Query对象
            return expandedQueries.stream()
                    .map(queryText -> query.metadata() == null
                            ? Query.from(queryText)
                            : Query.from(queryText, query.metadata()))
                    .collect(toList());
        } catch (Exception e) {
            log.error("查询扩展失败", e);
            return singletonList(query);
        }
    }

    /**
     * 格式化聊天历史记录
     */
    private String formatChatMemory(List<ChatMessage> chatMemory) {
        return chatMemory.stream()
                .map(this::formatChatMessage)
                .filter(Objects::nonNull)
                .collect(joining("\n"));
    }

    /**
     * 格式化单条聊天消息
     */
    private String formatChatMessage(ChatMessage message) {
        if (message instanceof UserMessage) {
            return "用户: " + ((UserMessage) message).singleText();
        } else if (message instanceof AiMessage aiMessage) {
            if (aiMessage.hasToolExecutionRequests()) {
                return null;
            }
            return "AI: " + aiMessage.text();
        } else {
            return null;
        }
    }

    /**
     * 解析响应中的查询行
     */
    private List<String> parseQueryLines(String response) {
        return Arrays.stream(response.split("\n"))
                .map(String::trim)
                .filter(line -> !line.isEmpty())
                .collect(toList());
    }

    /**
     * 查询转换策略
     */
    public enum TransformStrategy {
        /**
         * 原始查询，不做任何转换
         */
        ORIGINAL,

        /**
         * 压缩查询，将对话历史和当前查询压缩为一个简洁查询
         */
        COMPRESS,

        /**
         * 扩展查询，生成多个不同表述但含义相同的查询
         */
        EXPAND,

        /**
         * 混合策略，先压缩再扩展
         */
        HYBRID,

        /**
         * 智能策略，根据查询特点自动选择最合适的转换方式
         */
        AUTO
    }

    /**
     * 增强查询转换器构建器
     */
    public static class EnhancedQueryTransformerBuilder {
        private ChatModel ChatModel;
        private TransformStrategy strategy = TransformStrategy.AUTO;
        private PromptTemplate compressTemplate;
        private PromptTemplate expandTemplate;
        private int expansionCount = 3;
        private String systemPrompt;

        EnhancedQueryTransformerBuilder() {
        }

        public EnhancedQueryTransformerBuilder ChatModel(ChatModel ChatModel) {
            this.ChatModel = ChatModel;
            return this;
        }

        public EnhancedQueryTransformerBuilder strategy(TransformStrategy strategy) {
            this.strategy = strategy;
            return this;
        }

        public EnhancedQueryTransformerBuilder compressTemplate(PromptTemplate compressTemplate) {
            this.compressTemplate = compressTemplate;
            return this;
        }

        public EnhancedQueryTransformerBuilder expandTemplate(PromptTemplate expandTemplate) {
            this.expandTemplate = expandTemplate;
            return this;
        }

        public EnhancedQueryTransformerBuilder expansionCount(int expansionCount) {
            this.expansionCount = expansionCount;
            return this;
        }

        public EnhancedQueryTransformerBuilder systemPrompt(String systemPrompt) {
            this.systemPrompt = systemPrompt;
            return this;
        }

        public EnhancedQueryTransformerImpl build() {
            EnhancedQueryTransformerImpl transformer = new EnhancedQueryTransformerImpl(
                    this.ChatModel,
                    this.strategy,
                    this.compressTemplate,
                    this.expandTemplate,
                    this.expansionCount);

            if (this.systemPrompt != null) {
                transformer.setSystemPrompt(this.systemPrompt);
            }

            return transformer;
        }
    }
}